#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2023/2/10 16:44
# @Author  : YueJian
# @File    : main.py
# @Description :
import os
import os.path
import re
import shutil
import stat
from email.utils import formatdate
from mimetypes import guess_type
from pathlib import Path
from urllib.parse import quote
import aiofiles
import uvicorn
from fastapi import Body, File, Path as F_Path, UploadFile, Request
from fastapi import Form, FastAPI
from starlette.responses import StreamingResponse
from starlette.middleware.cors import CORSMiddleware
 
app = FastAPI(root_path="/api/uploader")
# 允许跨域
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
# 上传位置
UPLOAD_FILE_PATH = "./static/file"
 
 
@app.get("/file-slice")
async def check_file(identifier: str, totalChunks: int):
    """
    判断上传文件是否存在
    :param identifier:
    :param totalChunks:
    :return:
    """
    path = os.path.join(UPLOAD_FILE_PATH, identifier)
    if not os.path.exists(path):
        # 不存在
        return {"code": 200, "data": {"isExist": False}}
    else:
        # 存在时判断分片数是否
        chunks = [int(str(i).split("_")[1]) for i in os.listdir(path)]
        return {"code": 200, "data": {"isExist": False, "uploaded": chunks}}
 
 
@app.post("/file-slice")
async def upload_file(
        request: Request,
        identifier: str = Form(..., description="文件唯一标识符"),
        chunkNumber: str = Form(..., description="文件分片序号（初值为1）"),
        file: UploadFile = File(..., description="文件")
):
    """文件分片上传"""
    path = Path(UPLOAD_FILE_PATH, identifier)
    if not os.path.exists(path):
        os.makedirs(path)
    file_name = Path(path, f'{identifier}_{chunkNumber}')
    if not os.path.exists(file_name):
        async with aiofiles.open(file_name, 'wb') as f:
            await f.write(await file.read())
    return {"code": 200, "data": {
        'chunk': f'{identifier}_{chunkNumber}'
    }}
 
 
@app.put("/file-merge")
async def merge_file(
        request: Request,
        filename: str = Body(..., description="文件名称含后缀"),
        identifier: str = Body(..., description="文件唯一标识符"),
        totalChunks: int = Body(..., description="总分片数")
):
    """合并分片文件"""
    target_file_name = Path(UPLOAD_FILE_PATH, filename)
    path = Path(UPLOAD_FILE_PATH, identifier)
    try:
        async with aiofiles.open(target_file_name, 'wb+') as target_file:  # 打开目标文件
            for i in range(len(os.listdir(path))):
                temp_file_name = Path(path, f'{identifier}_{i + 1}')
                async with aiofiles.open(temp_file_name, 'rb') as temp_file:  # 按序打开每个分片
                    data = await temp_file.read()
                    await target_file.write(data)  # 分片内容写入目标文件
    except Exception as e:
        print(e)
        return {"code": 400, "msg": "合并失败"}
    shutil.rmtree(path)  # 删除临时目录
    return {"code": 200, "data": {"filename": filename, "url": f"/api/file/file-slice/{filename}"}}
 
 
@app.get("/file-slice/{file_name}")
async def download_file(request: Request, file_name: str = F_Path(..., description="文件名称（含后缀）")):
    """分片下载文件，支持断点续传"""
    # 检查文件是否存在
    file_path = Path(UPLOAD_FILE_PATH, file_name)
    if not os.path.exists(file_path):
        return {"code": 400, "msg": "文件不存在"}
    # 获取文件的信息
    stat_result = os.stat(file_path)
    content_type, encoding = guess_type(file_path)
    content_type = content_type or 'application/octet-stream'
    # 读取文件的起始位置和终止位置
    range_str = request.headers.get('range', '')
    range_match = re.search(r'bytes=(\d+)-(\d+)', range_str, re.S) or re.search(r'bytes=(\d+)-', range_str, re.S)
    if range_match:
        start_bytes = int(range_match.group(1))
        end_bytes = int(range_match.group(2)) if range_match.lastindex == 2 else stat_result.st_size - 1
    else:
        start_bytes = 0
        end_bytes = stat_result.st_size - 1
    # 这里 content_length 表示剩余待传输的文件字节长度
    content_length = stat_result.st_size - start_bytes if stat.S_ISREG(stat_result.st_mode) else stat_result.st_size
    # 构建文件名称
    name, *suffix = file_name.rsplit('.', 1)
    suffix = f'.{suffix[0]}' if suffix else ''
    filename = quote(f'{name}{suffix}')  # 文件名编码，防止中文名报错
    # 打开文件从起始位置开始分片读取文件
    return StreamingResponse(
        file_iterator(file_path, start_bytes, 1024 * 1024 * 5),  # 每次读取 5M
        media_type=content_type,
        headers={
            'content-disposition': f'attachment; filename="{filename}"',
            'accept-ranges': 'bytes',
            'connection': 'keep-alive',
            'content-length': str(content_length),
            'content-range': f'bytes {start_bytes}-{end_bytes}/{stat_result.st_size}',
            'last-modified': formatdate(stat_result.st_mtime, usegmt=True),
        },
        status_code=206 if start_bytes > 0 else 200
    )
 
 
def file_iterator(file_path, offset, chunk_size):
    """
    文件生成器
    :param file_path: 文件绝对路径
    :param offset: 文件读取的起始位置
    :param chunk_size: 文件读取的块大小
    :return: yield
    """
    with open(file_path, 'rb') as f:
        f.seek(offset, os.SEEK_SET)
        while True:
            data = f.read(chunk_size)
            if data:
                yield data
            else:
                break
 
 
if __name__ == '__main__':
    uvicorn.run(app=app, host="0.0.0.0", port=8000)